Skip to content

feat(pipeline): run pipeline stages in parallel - #873

Open
albertlast wants to merge 1 commit into
mayocream:mainfrom
albertlast:parallize
Open

feat(pipeline): run pipeline stages in parallel#873
albertlast wants to merge 1 commit into
mayocream:mainfrom
albertlast:parallize

Conversation

@albertlast

@albertlast albertlast commented Jul 25, 2026

Copy link
Copy Markdown

Motivation

I translate long manga — runs of 1000+ pages. At that scale the bottleneck isn't translation quality; the existing per-page quality is already good enough for what I need. It's wall-clock time. The pipeline currently runs one page through one stage at a time, so the GPU sits idle during OCR, the CPU sits idle during inference, and a long run takes far longer than the hardware requires.

This PR is aimed squarely at that: throughput on large batches, without changing what any single page produces.

What changed

Pages currently move through the pipeline strictly one stage at a time. This PR lets stages overlap, and lets a stage fold several pages into a single model call.

Two config knobs control it, both 0 = auto:

Knob Meaning 1 means
max_inflight_pages Pages moving through the pipeline at once Fully sequential — the escape hatch
max_batch_pages Pages a stage folds into one model call Stages still overlap, but no batching

max_batch_pages is the larger VRAM lever, so it's the one to lower first.

Which stage actually gets what

There are two independent wins here, and it's worth separating them:

1. Cross-stage overlap — applies to every stage, including the GPU-bound ones. Page 2 can run detection while page 1 is inpainting. No engine opts into this and none can opt out; it falls out of the streaming driver. For a 1000-page run this is the bulk of the gain.

2. Within-stage parallelism and batching — opt-in, per engine. The Engine trait gains max_workers, max_batch and run_batch, all defaulting to no concurrency. Only 4 of the 15 registered engines override them, and deliberately asymmetrically:

Stage Engine Workers Batch Why
Detection comic-text-detector (+ -seg, comic-text-bubble-detector, anime-text) 1 1 Shared GPU context; fanning out multiplies peak VRAM for no gain
OCR manga-ocr 1 ≤ 4 MangaOcr::inference is a genuine tensor batch — crops are cat'd into one forward pass
OCR paddle-ocr-vl-1.6, mit48px-ocr 1 1 Their batch API just loops internally, so batching would be a lie
Font detection yuzumarker-font-detection 1 ≤ 4 FontDetector::inference preprocesses crops in parallel then does one batched forward
Translation llm 4 (remote only) ≤ 4 (remote, and no custom system prompt) Remote providers are network-bound and stateless; a local llama.cpp context is &mut and would just queue on its state lock
Inpainting lama-manga, aot-inpainting, flux2-klein 1 1 Shared GPU context, as above
Rendering koharu-renderer min(cores, 4) 1 Pure CPU shaping/rasterisation on &self; the only shared state is a short-lived font-book mutex

run_batch returns one result per input, in order, so a failure is attributed to the page that caused it instead of failing the whole group.

Supporting changes:

  • Registry::get now dedupes concurrent misses behind a per-engine lock. Without it, parallel stages hitting a cold engine each load the model and allocate its GPU memory, then discard all but one.
  • Cross-page translation batching is skipped when a custom system prompt is in effect, since such a prompt describes the single-page [N] block format and can't be assumed to teach the batched [bP-N] form.

Stopping a run (and why stage threads are pooled)

Cancelling was where the naive version fell over, and it's what drove the least obvious part of this design.

The first cut gave each stage its own thread that exited when the stage ended. Pressing Stop ends every stage thread at once, and that turned out to crash the app in two separate ways:

  1. cuDNN teardown panics in a destructor. candle caches its cuDNN handles in a thread_local! (neither Send nor Sync), and cudarc's Drop for Cudnn unwraps cudnnDestroy. A thread that ends therefore destroys those handles and turns any teardown error into a panic inside a destructor — a hard crash, and cancelling is the easiest way to hit it.
  2. Shared hyper connections die with the runtime that opened them. The LLM client is one Arc<ClientWithMiddleware> shared by every stage. Dropping a stage's runtime takes its pooled connections with it, so a different, still-running worker's next request fails with error sending request.

So stage threads now park instead of exiting, and are reused across runs. That is the whole reason the pool exists — it isn't premature optimisation.

Cancellation semantics themselves:

  • Stop is checked once per batch at the top of each stage loop, so a cancelled run unwinds within one model call rather than mid-inference.
  • Each in-flight page carries its semaphore permit on the Item itself, so dropping an item anywhere — success, failure, or cancellation — releases its slot with no explicit bookkeeping and no leak path.
  • Every stage drains and exits on channel close, so stopping can't leave an orphaned worker behind.

There are unit tests for the pool covering exactly this: that thread-locals survive across stages, that a panicking stage does not kill its pooled thread, and that a closed-and-drained channel ends the stage cleanly.

Known gap: MCP does not honour these settings

start_pipeline (the HTTP route the UI uses) reads the limits from app config, so the Settings knobs apply there. The MCP entry point still passes PipelineLimits::default() — i.e. auto — so a whole-project run driven over MCP will parallelise even if the user set max_inflight_pages = 1. The pipeline CLI binary does the same, though it's single-page so it makes no practical difference there.

I left it as-is rather than guess at the intended layering: these limits are currently app config, not part of StartPipelineRequest, so it isn't obvious whether MCP should read global config or whether the limits belong on the request instead (which is roughly what #830 does with its per-run flag). Happy to wire it up either way — just say which you'd prefer.

Auto mode is not a black box

0 means auto, and the Settings UI resolves and displays what auto actually picked on this machine rather than leaving the user guessing:

  • Pages In Flight, when auto → "Auto: 5 — one page per pipeline step."
  • Batch Size, when auto → "Auto: each engine decides its own batch."
  • Always shown → "CPU-bound stages use 16 workers on this machine."

The worker count depends on the host's core count, which the frontend can't know, so MetaInfo gains a cpu_workers field reported by the server (GET /meta). The effective max_inflight_pages / max_batch_pages are exposed there too, so the displayed numbers are the ones the driver will really use — not a client-side guess.

⚠️ Likely conflicts with #830 (chapter context translation)

I want to flag this up front rather than have a maintainer discover it at merge time.

Page completion order becomes non-deterministic. That is inherent to this PR, not incidental: pages are streamed through stages and finish out of order, which is exactly why a Tracker::frontier exists — progress is reported from the lowest unfinished page so the percentage stays monotonic even when page 3 finishes before page 1. There are tests asserting this (frontier_never_decreases_under_interleaved_completion).

#830 wants the opposite for its chapter path. Its own settings string says it will "run detect and OCR on all pages first, then translate with shared chapter context", and it advertises that "translation preserves reading order across pages." That is a barrier plus a stable global ordering — the two properties this PR deliberately gives up on the default path.

Structurally the two collide in the same place. #830 splits run() into run_sequential() + run_chapter_mode() and inserts at @@ -149,6 +193,66 @@ pub async fn run(. This PR replaces the body of that same run() with the streaming stage-worker driver. We also both touch pipeline/engine.rs, llm.rs, bin/pipeline.rs, koharu-llm/src/prompt.rs, rpc/mcp/mod.rs, rpc/routes/pipelines.rs, SettingsDialog.tsx, openapi.json and the locale files. Whichever lands second will need a real rebase, not a mechanical one.

That said, I don't think they're fundamentally incompatible — and the two are arguably chasing different goals. #830 trades time for cross-page quality; this PR trades ordering for time. Its chapter mode is an early return into its own barriered code path, so it would keep its ordering guarantees regardless of what the default path does — the parallel driver simply wouldn't apply to it. The reconciliation is mostly mechanical-but-tedious rather than a design dead end, and run_step_for_all_pages could later use the same batching machinery.

User-visible behavior

  • Multi-page runs get faster; single-page runs are unchanged.
  • Pages finish out of order. Progress percentage stays monotonic, but per-page completion events no longer arrive in page order.
  • New Parallel Processing section in Settings for both knobs, translated into all nine locales. Saved immediately and applied to the next run — no restart.
  • Auto mode reports the concrete numbers it resolved to (see above), so 0 is inspectable rather than opaque.
  • Defaults are auto, so existing configs need no change. Old config files without the new fields still parse.
  • Setting max_inflight_pages = 1 restores the previous sequential behavior exactly, ordering included.

How I verified

  • cargo fmt -- --check, cargo check, cargo clippy -- -D warnings — clean.
  • cargo test --workspace --tests — all pass.
  • bun lint:ui and bun run --filter ui test (156 tests, 23 files) — all pass.
  • bun run generate:openapi reproduces the committed ui/openapi.json with no drift.
  • New unit tests cover batch draining, channel close/drain shutdown, progress-frontier monotonicity under interleaved completion, percentage clamping, config patch/round-trip (including 0 as a real "auto" value), and the thread pool — that thread-locals survive across stages, that a panicking stage doesn't kill its pooled thread, and that concurrent stages each get their own thread.
  • Beyond the test suite, this is what I actually run my own multi-hundred-page translations on.

I did not regenerate tests/integration-tests/client. It's already out of sync with main's own openapi.json and regenerating pulls in 250+ unrelated files from a different generator version, which would bury this diff. Happy to do it in a separate PR if you'd like.

AI usage disclosure

Per the contributing guide: AI assistance (Claude Code) was used for parts of this patch. I have reviewed, run and understood the change, and I'm responsible for it. Happy to walk through any part of the design — particularly the thread-pool parking rationale and the ordering trade-off above, which are the least obvious pieces.

🤖 Generated with Claude Code

Pages were processed strictly one stage at a time. This lets stages
overlap and lets a stage fold several pages into one model call.

Two config knobs control it, both `0` = auto:

- `max_inflight_pages` — pages moving through the pipeline at once.
  Set to `1` to restore fully sequential behaviour.
- `max_batch_pages` — pages a single stage folds into one model call.
  The larger VRAM lever, so lower this one first.

Engines opt in rather than out. The `Engine` trait gains `max_workers`,
`max_batch` and `run_batch`, all defaulting to no concurrency, so an
engine that shares one GPU context or a `&mut` model is untouched.
`run_batch` returns one result per input in order, so a failure is
attributed to the page that caused it instead of failing the group.

Supporting changes:

- `Registry::get` dedupes concurrent misses behind a per-engine lock, so
  parallel stages hitting a cold engine no longer each load the model
  and allocate its GPU memory.
- Stage threads are pooled and park between runs instead of exiting:
  candle caches cuDNN handles in a `thread_local!` that is unsafe to
  tear down.
- Translation batching is skipped when a custom system prompt is set,
  since such a prompt describes the single-page `[N]` block format.
- Settings UI for both knobs, translated into all nine locales.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

Thanks for your first PR to Koharu.

Please review our contribution guide before review:
https://koharu.rs/contribute/introduction/

In the PR description, include:

  • what changed
  • any user-visible behavior differences
  • how you verified the change

If AI helped produce the patch, a human still needs to review and understand it before submission.

@github-actions github-actions Bot added dependencies Pull requests that update a dependency file area: ui area: rpc labels Jul 25, 2026
@albertlast

Copy link
Copy Markdown
Author

when i see correctly a pr with similiar scoping got merged, so can i close this pr? @mayocream

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area: rpc area: ui dependencies Pull requests that update a dependency file

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant